{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deploy and Distribute TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook you will learn how to deploy TensorFlow models to TensorFlow Serving (TFS), using the REST API or the gRPC API, and how to train a model across multiple devices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import sklearn\n",
    "import sys\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"python\", sys.version)\n",
    "for module in mpl, np, pd, sklearn, tf, keras:\n",
    "    print(module.__name__, module.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert sys.version_info >= (3, 5) # Python ≥3.5 required\n",
    "assert tf.__version__ >= \"2.0\"    # TensorFlow ≥2.0 required"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Exercise](https://c1.staticflickr.com/9/8101/8553474140_c50cf08708_b.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1 – Deploying a Model to TensorFlow Serving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save/Load a `SavedModel`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(X_train_full, y_train_full), (X_test, y_test) = keras.datasets.fashion_mnist.load_data()\n",
    "X_train_full = X_train_full / 255.\n",
    "X_test = X_test / 255.\n",
    "X_valid, X_train = X_train_full[:5000], X_train_full[5000:]\n",
    "y_valid, y_train = y_train_full[:5000], y_train_full[5000:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = keras.models.Sequential([\n",
    "    keras.layers.Flatten(input_shape=[28, 28]),\n",
    "    keras.layers.Dense(100, activation=\"relu\"),\n",
    "    keras.layers.Dense(10, activation=\"softmax\")\n",
    "])\n",
    "model.compile(loss=\"sparse_categorical_crossentropy\",\n",
    "              optimizer=keras.optimizers.SGD(lr=1e-3),\n",
    "              metrics=[\"accuracy\"])\n",
    "model.fit(X_train, y_train, epochs=10, validation_data=(X_valid, y_valid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"my_fashion_mnist\"\n",
    "!rm -rf {MODEL_NAME}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "model_version = int(time.time())\n",
    "model_path = os.path.join(MODEL_NAME, str(model_version))\n",
    "os.makedirs(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.saved_model.save(model, model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for root, dirs, files in os.walk(MODEL_NAME):\n",
    "    indent = '    ' * root.count(os.sep)\n",
    "    print('{}{}/'.format(indent, os.path.basename(root)))\n",
    "    for filename in files:\n",
    "        print('{}{}'.format(indent + '    ', filename))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli show --dir {model_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli show --dir {model_path} --tag_set serve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli show --dir {model_path} --tag_set serve \\\n",
    "                      --signature_def serving_default"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli show --dir {model_path} --all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Warning**: if you are using an old version of TensorFlow, the method name for the `serving_default` signature may be empty, due to [TensorFlow issue #25235](https://github.com/tensorflow/tensorflow/issues/25235). In this case, you need to use `keras.experimental.export()` instead of `tf.saved_model.save()` (or upgrade to the latest TensorFlow version)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's write a few test instances to a `npy` file so we can pass them easily to our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_new = X_test[:3]\n",
    "np.save(\"my_fashion_mnist_tests.npy\", X_new, allow_pickle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_name = model.input_names[0]\n",
    "input_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now let's use `saved_model_cli` to make predictions for the instances we just saved:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli run --dir {model_path} --tag_set serve \\\n",
    "                     --signature_def serving_default    \\\n",
    "                     --inputs {input_name}=my_fashion_mnist_tests.npy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TensorFlow Serving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install [Docker](https://docs.docker.com/install/) if you don't have it already. Then run:\n",
    "\n",
    "```bash\n",
    "docker pull tensorflow/serving\n",
    "\n",
    "docker run -it --rm -p 8501:8501 \\\n",
    "   -v \"`pwd`/my_fashion_mnist:/models/my_fashion_mnist\" \\\n",
    "   -e MODEL_NAME=my_fashion_mnist \\\n",
    "   tensorflow/serving\n",
    "```\n",
    "\n",
    "Once you are finished using it, press Ctrl-C to shut down the server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "input_data_json = json.dumps({\n",
    "    \"signature_name\": \"serving_default\",\n",
    "    \"instances\": X_new.tolist(),\n",
    "})\n",
    "print(input_data_json[:200] + \"...\" + input_data_json[-200:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's use TensorFlow Serving's REST API to make predictions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "SERVER_URL = 'http://localhost:8501/v1/models/my_fashion_mnist:predict'\n",
    "            \n",
    "response = requests.post(SERVER_URL, data=input_data_json)\n",
    "response.raise_for_status()\n",
    "response = response.json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_proba = np.array(response[\"predictions\"])\n",
    "y_proba.round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Serialized Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "serialized = []\n",
    "for image in X_new:\n",
    "    image_data = tf.train.FloatList(value=image.ravel())\n",
    "    features = tf.train.Features(\n",
    "        feature={\n",
    "            \"image\": tf.train.Feature(float_list=image_data),\n",
    "        }\n",
    "    )\n",
    "    example = tf.train.Example(features=features)\n",
    "    serialized.append(example.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[data[:100]+b'...' for data in serialized]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_images(serialized):\n",
    "    expected_features = {\n",
    "        \"image\": tf.io.FixedLenFeature([28 * 28], dtype=tf.float32)\n",
    "    }\n",
    "    examples = tf.io.parse_example(serialized, expected_features)\n",
    "    return tf.reshape(examples[\"image\"], (-1, 28, 28))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parse_images(serialized)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "serialized_inputs = keras.layers.Input(shape=[], dtype=tf.string)\n",
    "images = keras.layers.Lambda(lambda serialized: parse_images(serialized))(serialized_inputs)\n",
    "y_proba = model(images)\n",
    "ser_model = keras.models.Model(inputs=[serialized_inputs], outputs=[y_proba])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SER_MODEL_NAME = \"my_ser_fashion_mnist\"\n",
    "\n",
    "ser_model_version = int(time.time())\n",
    "ser_model_path = os.path.join(SER_MODEL_NAME, str(ser_model_version))\n",
    "os.makedirs(ser_model_path)\n",
    "tf.saved_model.save(ser_model, ser_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli show --dir {ser_model_path} --all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```bash\n",
    "docker run -it --rm -p 8500:8500 -p 8501:8501 \\\n",
    "   -v \"`pwd`/my_ser_fashion_mnist:/models/my_ser_fashion_mnist\" \\\n",
    "   -e MODEL_NAME=my_ser_fashion_mnist \\\n",
    "   tensorflow/serving\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "import json\n",
    "\n",
    "ser_input_data_json = json.dumps({\n",
    "    \"signature_name\": \"serving_default\",\n",
    "    \"instances\": [{\"b64\": base64.b64encode(data).decode(\"utf-8\")}\n",
    "                  for data in serialized],\n",
    "})\n",
    "print(ser_input_data_json[:200] + \"...\" + ser_input_data_json[-200:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "SER_SERVER_URL = 'http://localhost:8501/v1/models/my_ser_fashion_mnist:predict'\n",
    "            \n",
    "response = requests.post(SER_SERVER_URL, data=ser_input_data_json)\n",
    "response.raise_for_status()\n",
    "response = response.json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_proba = np.array(response[\"predictions\"])\n",
    "y_proba.round(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pip install --no-deps tensorflow-serving-api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import grpc\n",
    "from tensorflow_serving.apis import predict_pb2\n",
    "from tensorflow_serving.apis import prediction_service_pb2_grpc\n",
    "\n",
    "channel = grpc.insecure_channel('localhost:8500')\n",
    "predict_service = prediction_service_pb2_grpc.PredictionServiceStub(channel)\n",
    "\n",
    "request = predict_pb2.PredictRequest()\n",
    "request.model_spec.name = SER_MODEL_NAME\n",
    "request.model_spec.signature_name = \"serving_default\"\n",
    "input_name = ser_model.input_names[0]\n",
    "request.inputs[input_name].CopyFrom(tf.compat.v1.make_tensor_proto(serialized))\n",
    "\n",
    "result = predict_service.Predict(request, 10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_name = ser_model.output_names[0]\n",
    "output_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shape = [dim.size for dim in result.outputs[output_name].tensor_shape.dim]\n",
    "shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_proba = np.array(result.outputs[output_name].float_val).reshape(shape)\n",
    "y_proba.round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Exercise](https://c1.staticflickr.com/9/8101/8553474140_c50cf08708_b.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2 – Distributed Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "keras.backend.clear_session()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distribution = tf.distribute.MirroredStrategy()\n",
    "\n",
    "with distribution.scope():\n",
    "    model = keras.models.Sequential([\n",
    "        keras.layers.Flatten(input_shape=[28, 28]),\n",
    "        keras.layers.Dense(100, activation=\"relu\"),\n",
    "        keras.layers.Dense(10, activation=\"softmax\")\n",
    "    ])\n",
    "    model.compile(loss=\"sparse_categorical_crossentropy\",\n",
    "                  optimizer=keras.optimizers.SGD(lr=1e-3),\n",
    "                  metrics=[\"accuracy\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(X_train, y_train, epochs=10, validation_data=(X_valid, y_valid), batch_size=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
